from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, TextLoader  # 替换为 DirectoryLoader
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaEmbeddings, ChatOllama
from langserve import add_routes
from fastapi import FastAPI
import uvicorn

# 1. 使用 DirectoryLoader 加载目录中的所有文件
loader = DirectoryLoader(
    path='./txt',  # 替换为你的目录路径
    glob="**/*.txt",  # 匹配所有 .txt 文件，可以修改为其他文件类型
    loader_cls=TextLoader,  # 指定使用 TextLoader 加载文件
    loader_kwargs={"encoding": "utf-8"},  # 传递 encoding 参数
    recursive=True,  # 递归加载子目录中的文件
)
documents = loader.load()

# 2. 将文档分割成小块
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)

# 3. 创建嵌入并向量存储
embeddings = OllamaEmbeddings(model="nomic-embed-text:latest")
docsearch = FAISS.from_documents(texts, embeddings)
retriever = docsearch.as_retriever(search_kwargs={"k": 4})  # 检索前 4 个相关文档

llm = ChatOllama(
    base_url="http://localhost:11434",
    model="deepseek-r1:1.5b",
    stream=True
)

# ChatPromptTemplate 定义
prompt_template = ChatPromptTemplate.from_messages([
    ("system", "请直接根据文档内容回答问题，不要进行额外的推理或生成。已知信息：{context}"),
    ("human", "{question}")
])

# 自定义的提示词参数
chain_type_kwargs = {
    "prompt": prompt_template,
}

# 定义RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",  # 使用stuff模式将上下文拼接到提示词中
    chain_type_kwargs=chain_type_kwargs,
    retriever=retriever
)

# 6. 创建 FastAPI 应用并添加路由
app = FastAPI()
add_routes(
    app,
    qa_chain,
    path="/topic",
)

if __name__ == "__main__":
    uvicorn.run(app, host="localhost", port=8000)
